import os
import pandas as pd
import collections


# 处理原始数据，将评分数据存储
def generateUserRatingData(ratingPath, itemPath, savetoPath):
    print('开始处理数据...')
    # 评分数据
    df_rating = pd.read_csv(open(ratingPath, 'r', encoding='utf-8'), sep='\t', header=None)

    # 电影列表
    item_info = pd.read_csv(open(itemPath, 'r', encoding='ISO-8859-1'), sep='|', header=None)

    movie_year = [str(item).split("-")[-1] for item in item_info[2].dropna()]

    # 拿到item数据的第一列，电影名
    itemList = [str(item_info[1].tolist()[index]) + ';' + str(index + 1) for index in range(len(item_info[1].tolist()))]

    # item数据大小（多少个电影）
    itemSize = len(itemList)

    # user数据大小（多少个user）
    userSize = len(df_rating[0].drop_duplicates().tolist())

    """
    计算评分数据中的每个item被评分的次数
    item_counts：{电影id1: 评论次数, 电影id2: 评论次数,电影id3: 评论次数, ... , 电影id1682: 评论次数 }
    item_counts示例: {50: 583, 258: 509, 100: 508, 181: 507, 294: 485, 286: 481, 288: 478,...}
    """
    item_counts = collections.Counter(list(df_rating[1]))

    # 用户评分矩阵列名构造，增加user列，结果形式是：用户id, 电影id1, 电影id2, 电影id3, ..., 电影id1682
    df_write = pd.DataFrame(columns=['user'] + itemList)

    # 去除少于50个用户评分的item
    min_ratings = 50
    # 需要删除的item
    removieItem = []
    for i in range(1, userSize + 1):
        tmpItem = [0 for j in range(itemSize)]
        # 获取user id==i的用户的数据
        """
               0    1  2          3
        202    1   61  4  878542420
        305    1  189  3  888732928
        333    1   33  4  878542699
        334    1  160  4  875072547
        478    1   20  4  887431883
        ...   ..  ... ..        ...
        92049  1   28  4  875072173
        92487  1  172  5  874965478
        94019  1  122  3  875241498
        96699  1  152  5  878542589
        99073  1   94  2  875072956
        """
        df_tmp = df_rating[df_rating[0] == i]
        for k in df_tmp.index:
            # 获取大于50个用户评分数的item
            if item_counts[df_tmp.loc[k][1]] >= min_ratings:
                tmpItem[df_tmp.loc[k][1] - 1] = df_tmp.loc[k][2]  # tmpItem存的是评分，下标是item序列-1
            else:
                removieItem.append(df_tmp.loc[k][1])  # 需要删除的item对应的数字序列
        """
        tmpItem内容
        [5, 3, 4, 3, 3, 0, 4, 1, 5, 3, 2, 5, 5, 5, 5, 0, 3, 0, 5, 4, 1, 4, 4, 3, 4, 3, 2, 4, 1, 0, 3, 5, 4, 0, 0, 0, 0, 3, 4, 3, 0, 5, 0, 5, 5, 0, 4, 5, 3, 5, 4, 4, 3, 3, 5, 4, 0, 4, 5, 5, 4, 3, 2, 5, 4, 4, 3, 4, 3, 3, 3, 4, 3, 0, 0, 4, 4, 0, 4, 4, 5, 5, 3, 0, 3, 5, 5, 4, 5, 4, 5, 3, 5, 2, 4, 5, 3, 4, 3, 5, 2, 2, 0, 0, 2, 4, 0, 5, 5, 0, 5, 0, 0, 5, 0, 3, 3, 3, 0, 1, 4, 3, 4, 5, 3, 2, 5, 4, 5, 0, 1, 4, 4, 4, 4, 3, 5, 0, 3, 1, 3, 2, 1, 4, 2, 0, 3, 2, 0, 5, 4, 5, 3, 5, 2, 4, 4, 3, 3, 4, 4, 4, 4, 3, 5, 5, 2, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 3, 3, 5, 4, 5, 4, 4, 4, 4, 3, 3, 5, 5, 4, 4, 4, 5, 5, 5, 5, 4, 3, 3, 5, 4, 5, 3, 4, 5, 5, 4, 4, 3, 4, 2, 4, 3, 5, 3, 3, 1, 3, 5, 4, 5, 0, 2, 3, 4, 5, 4, 4, 1, 3, 2, 4, 5, 0, 2, 4, 4, 3, 4, 5, 1, 0, 2, 5, 0, 4, 4, 4, 0, 2, 0, 1, 2, 0, 4, 5, 1, 1, 0, 3, 0, 2, 4, 0, 0, 5, 5, 5, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
        """
        df_write.loc[i] = [i] + tmpItem  # 一行数据，用户Id，及其对item的评分

    removieItem = list(set(removieItem))
    df_write.drop(df_write.columns[removieItem], axis=1, inplace=True)  # inplace=True 直接替换原数组
    df_write.to_csv(savetoPath, index=None)  # 处理后的数据存入文件中
    return item_info, df_write  # item_info存储的是关于item的内容属性,df_write存储的是经预处理的评分数据


# 读取和存储关于item内容
def generateUserItemContent(item_info, df_write, savetoPath):
    # 获得项目名对应的数字
    itemList = [int(itemNum.split(';')[-1]) for itemNum in
                df_write.columns[1:]]  # 获得所有列名对应的数字,df_write中的第0列是人名，从第1列开始是项目名
    # 这是项目从属的类型，比如未知，动作，冒险等
    itemKind = ['unknown', 'Action', 'Adventure', 'Animation', 'Children\'s',
                'Comedy', 'Crime', 'Documentary', 'Drama', 'Fantasy',
                'Film-Noir', 'Horror', 'Musical', 'Mystery', 'Romance',
                'Sci-Fi', 'Thriller', 'War', 'Western']
    item_kind = pd.DataFrame(columns=[['item_id'] + itemKind])
    start = 5
    count = 0
    for itemNum in itemList:  # 列名对应的数字
        item_kind.loc[count] = [itemNum] + item_info.iloc[itemNum - 1][start:].tolist()
        count += 1
    item_kind.to_csv(savetoPath, index=None)


if __name__ == '__main__':
    data_path = os.path.dirname(os.path.dirname(os.getcwd()))
    # 评分矩阵
    ratingMatrixPath = os.path.join(data_path, r'data\ml-100k\u.data')  # 用户对电影的评分数据：用户id、电影id、评分、时间戳
    moviesPath = os.path.join(data_path, r'data\ml-100k\u.item')  # 电影的基础信息：电影id、电影名称、电影上映时间、电影发布时间、IMDb URL、电影类型
    userMoviesRatingSavePath = os.path.join(data_path,
                                            r'data\processedData\ratingMatrix.csv')  # 电影的标签数据: 用户id、电影id1,电影id2,电影id3,...,电影id1682
    moviesContentSavePath = os.path.join(data_path,
                                         r'data\processedData\moviesContent.csv')  # 电影的标签信息：电影id、 ["未知", "动作", "冒险", "动画", "儿童", "喜剧", "犯罪", "纪录片", "剧情", "奇幻", "黑色", "恐怖", "音乐剧", "悬疑", "浪漫", "科幻", "惊悚", "战争", "西部"]
    item_info, df_write = generateUserRatingData(ratingMatrixPath, moviesPath, userMoviesRatingSavePath)
    generateUserItemContent(item_info, df_write, moviesContentSavePath)
